"""
本代码是.py的gpt改编版.添加了代码规范,使得更容易阅读.
基于计算机视觉的颜色检测与道路跟踪系统
功能：
1. 通过摄像头检测特定颜色区域
2. 根据颜色区域面积判断不同场景
3. 通过串口发送控制指令
4. 实时显示处理后的图像
"""

import cv2
import numpy as np
import serial
from TimeStatusKeeper import TimeStatusKeeper
import threading
from detect import RoadDetector  # 自定义道路检测模块


COLOR_THRESHOLDS = {
    "purple": {"lower": np.array([124, 77, 51]), "upper": np.array([177, 255, 255])},
    "brown":  {"lower": np.array([0, 96, 0]),   "upper": np.array([11, 255, 108])},
    "yellow": {"lower": np.array([21, 80, 128]), "upper": np.array([50, 255, 255])},
    "blue":   {"lower": np.array([100, 43, 46]), "upper": np.array([124, 255, 255])},
    "green":  {"lower": np.array([56, 46, 89]),  "upper": np.array([89, 255, 255])},
    "red":    {"lower": np.array([0, 90, 163]),  "upper": np.array([10, 255, 255])},
    "orange": {"lower": np.array([8, 128, 117]), "upper": np.array([12, 255, 160])},
}

# | 颜色 | HSV 值（°，%）    |                            |
# | 紫色 | 270°, 100%, 100% |                            |
# | 棕色 | 30°, 59%, 41%    |                            |
# | 黄色 | 60°, 100%, 100%  |                            |
# | 蓝色 | 240°, 100%, 100% |                            |
# | 绿色 | 120°, 100%, 100% |                            |
# | 红色 | 0°, 100%, 100%   |                            |
# | 橘色 | 30°, 100%, 100%  | ([Medium][1], [Reddit][2]) |

STATE_HOME = 0    # 初始/返回状态
STATE_TURNING_LFFT = 123 # 转弯状态
STATE_TURNING_RIGHT = 321
GO_HOME = 666
current_state = STATE_HOME  # 当前系统状态（原rightcnt）
time_status = 0  # 计时状态：0-停止，1-计时中，2-计时完成
stop_event = threading.Event()  # 用于控制线程停止
blueFlag = False # 有时候会将电线识别成蓝色横条，需要加上flag限定 
ser = serial.Serial("/dev/ttyUSB0", 115200)  # 串口设备
keeper = TimeStatusKeeper()
mouse_data = {'hsv': (0, 0, 0), 'pos': (0, 0), 'frame': None}

SERIAL_CMD = {
    'TURN':      'B\r\n',
    'HOME':      'T\r\n',
    'RIGHT_BALL':'P\r\n',
    'LEFT_BALL': 'Z\r\n',
    'STAIR':     'A\r\n'
}

    
def calculate_max_color_area(frame, lower_bound, upper_bound):
    """
    计算指定颜色区域的最大连通面积
    :param frame: 输入BGR图像帧
    :param lower_bound: HSV下限
    :param upper_bound: HSV上限
    :return: 最大连通区域面积（像素数）
    """
    max_area = 0
    
    # 图像预处理
    hsv_img = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    blurred_img = cv2.GaussianBlur(hsv_img, (5, 5), 0)
    
    # 颜色阈值处理
    color_mask = cv2.inRange(blurred_img, lower_bound, upper_bound)
    
    # 仅关注图像下半部分（ROI区域）
    height, width = color_mask.shape
    roi_mask = color_mask[height//2 : height, :]
    
    # 形态学操作（腐蚀去噪）
    kernel = np.ones((3, 3), np.uint8)
    eroded_mask = cv2.erode(roi_mask, kernel, iterations=2)
    
    # 寻找轮廓并计算最大面积
    contours, _ = cv2.findContours(
        eroded_mask.copy(), 
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )
    
    for contour in contours:
        area = cv2.contourArea(contour)
        max_area = max(max_area, area)
    
    return max_area

def determine_scenario(color_areas,movementFlag):
    """
    根据颜色区域面积判断当前场景并发送控制指令
    :param color_areas: 各颜色区域面积数组（索引对应初始化顺序）
    """
    global current_state
    
    ORANGE, BLUE, RED, GREEN, BROWN, PURPLE = 0, 1, 2, 3, 4, 5
    


    # 场景判断逻辑
    if (color_areas[RED] + color_areas[GREEN] > 20000 
        and color_areas[RED] > 10000 
        and color_areas[GREEN] > 10000
        and keeper.status == STATE_HOME):
        print("[场景] 检测到转弯区域")
        keeper.set_status(10)
        print(f"send {SERIAL_CMD['TURN']}")
        movementFlag = False
        ser.write(SERIAL_CMD['TURN'].encode())
    
    elif (color_areas[ORANGE] > 20000 
          and color_areas[GREEN] < 15000   
          and keeper.status == GO_HOME):
        print("[场景] 返回基地")
        keeper.set_status(10)
        print(f"send {SERIAL_CMD['HOME']}")
        movementFlag = False
        ser.write(SERIAL_CMD['HOME'].encode())
    
    elif (color_areas[PURPLE] > 50000
        and keeper.status == STATE_TURNING_RIGHT):
        print("[场景] 检测到右侧球体")
        keeper.set_status(10)
        print(f"send {SERIAL_CMD['RIGHT_BALL']}")
        movementFlag = False
        ser.write(SERIAL_CMD['RIGHT_BALL'].encode())
    
    elif (color_areas[BROWN] > 7500    
          and color_areas[GREEN] < 50000 and keeper.status == STATE_TURNING_LFFT):
        print("[场景] 检测到左侧球体")
        keeper.set_status(10)
        print(f"send {SERIAL_CMD['LEFT_BALL']}")
        ser.write(SERIAL_CMD['LEFT_BALL'].encode())
    
    elif color_areas[BLUE] > 25000 and keeper.status == STATE_HOME:
        print("[场景] 检测到楼梯")
        print(f"send {SERIAL_CMD['STAIR']}")
        movementFlag = False
        ser.write(SERIAL_CMD['STAIR'].encode())
        
def show_hsv_on_hover(event, x, y, flags, param):
    """鼠标回调函数，用于更新HSV值和坐标"""
    if event == cv2.EVENT_MOUSEMOVE:
        # 将当前鼠标位置保存到共享字典中
        param['pos'] = (x, y)
        # 从共享字典中获取当前帧
        frame = param.get('frame')
        # 确保帧存在且坐标在范围内
        if frame is not None and 0 <= y < frame.shape[0] and 0 <= x < frame.shape[1]:
            # 将BGR像素转换为HSV
            bgr_pixel = frame[y, x]
            hsv_pixel = cv2.cvtColor(np.uint8([[bgr_pixel]]), cv2.COLOR_BGR2HSV)[0][0]
            param['hsv'] = (hsv_pixel[0], hsv_pixel[1], hsv_pixel[2])

def main():
    global current_state, time_status
    # 初始化摄像头
    camera = cv2.VideoCapture(0)

    # 初始化道路检测器
    detector = RoadDetector(
        roadThreshold=[(np.array([0,0,95]),np.array( [255,255,255]))],
        stairThreshold=(np.array([69,0,131]),np.array ([121,113,255])),
        backThreshold=(np.array([83,94,79]),np.array([104,175,207])),
        maskKSize=3
    )

   
    while True:
        movementFlag = True   # 是否发送move_x rotation 
        # 多帧采样取平均（降噪）
        color_areas = np.zeros(6)
        for _ in range(5):
            ret, frame = camera.read()
            if not ret: continue
            
            # 计算各颜色区域面积（保持初始化顺序）
            color_areas[0] += calculate_max_color_area(frame, COLOR_THRESHOLDS["orange"]["lower"], COLOR_THRESHOLDS["orange"]["upper"])
            color_areas[1] += calculate_max_color_area(frame, COLOR_THRESHOLDS["blue"]["lower"], COLOR_THRESHOLDS["blue"]["upper"])
            color_areas[2] += calculate_max_color_area(frame, COLOR_THRESHOLDS["red"]["lower"], COLOR_THRESHOLDS["red"]["upper"])
            color_areas[3] += calculate_max_color_area(frame, COLOR_THRESHOLDS["green"]["lower"], COLOR_THRESHOLDS["green"]["upper"])
            color_areas[4] += calculate_max_color_area(frame, COLOR_THRESHOLDS["brown"]["lower"], COLOR_THRESHOLDS["brown"]["upper"])
            color_areas[5] += calculate_max_color_area(frame, COLOR_THRESHOLDS["purple"]["lower"], COLOR_THRESHOLDS["purple"]["upper"])

        print("各颜色区域面积:", color_areas)
        determine_scenario(color_areas,movementFlag)

        # 图像预处理（道路检测）
        processed_img = process_image_for_tracking(frame, keeper.status)
        
        # 更新道路检测器
        detector.update(processed_img) #在这里计算斜率 楼梯质心 楼梯顶部
        tracking_data = detector.getTrackingData()
        
        # 运动控制逻辑
        if tracking_data is not None and movementFlag is True:
            handle_movement_control(tracking_data)
        
        # cv2.imshow('test',processed_img)
        show_debug_windows(processed_img, detector, frame)

        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    
    camera.release()
    cv2.destroyAllWindows()

def process_image_for_tracking(original_img, state):
    """
    道路检测图像预处理流程
    :param original_img: 原始图像帧
    :param state: 当前系统状态
    :return: 处理后的HSV图像
    """
    resized_img = cv2.resize(
        original_img, 
        (int(original_img.shape[1]//3 * 2), int(original_img.shape[0]//3 * 2))
    )
     
    blurred_img = cv2.GaussianBlur(resized_img,(5,5),0) 
    # ROI区域截取（根据不同状态调整）
    if state == STATE_TURNING_LFFT or state == STATE_TURNING_RIGHT:
        roi_img = blurred_img[blurred_img.shape[0]//3+110 : blurred_img.shape[0], 
                           0 : blurred_img.shape[1]]
    else:
        roi_img = blurred_img[blurred_img.shape[0]//3+110 : blurred_img.shape[0], :]
    
    hsv_img=cv2.cvtColor(roi_img,cv2.COLOR_BGR2HSV)

    h,s,v=cv2.split(hsv_img)
    clahe = cv2.createCLAHE(clipLimit=2.0,tileGridSize=(8,8))
    v=clahe.apply(v)
    hsv_img = cv2.merge([h,s,v])
    enhanced_display = cv2.cvtColor(hsv_img,cv2.COLOR_HSV2BGR)
    # cv2.imshow("test",enhanced_display)
    return hsv_img

def handle_movement_control(tracking_data):
    """
    处理运动控制逻辑
    :param tracking_data: 道路检测器返回的跟踪数据
    """
    slope, intercept, start_point, end_point = tracking_data
    
    # 计算中线偏移
    mid_offset = ((start_point[0] + end_point[0])/2 - 160) / 320
    
    # 限制斜率范围
    slope = np.clip(slope, -1.5, 1.2)
    #print(f'slope: {slope}')
    #print(f'mid_offset: {mid_offset}')
    
    # 计算控制参数
    move_x = round(mid_offset * 100 / 2.5) / 100
    rotation = round((-slope*10*100/180 + mid_offset*10)*1.3*100)/100 
    
    # 发送控制指令
    print (f'move : {move_x} , rotation: {rotation} ')
    n = ser.write(f"{move_x} {rotation} \r\n".encode())
    if n == len(f"{move_x} {rotation} \r\n"):
        pass
    else:
        print(f"仅发送了 {n} 字节，预期 {len(2)}")

def show_debug_windows(processed_img, detector, original_frame):
    """
    显示调试窗口
    :param processed_img: 处理后的图像
    :param detector: 道路检测器实例
    :param original_frame: 原始图像帧
    """
    # print("processed_img 类型:", type(processed_img), "形状:", processed_img.shape if hasattr(processed_img, 'shape') else "N/A")
    # print("detector.img_mask 类型:", type(detector.img_mask), "形状:", detector.img_mask.shape if hasattr(detector.img_mask, 'shape') else "N/A")
    # print("original_frame 类型:", type(original_frame), "形状:",  _frame.shape if hasattr(original_frame, 'shape') else "N/A")

    detector.showTracking(processed_img)
    cv2.imshow("Road Mask", detector.img_mask)
    # cv2.imshow("Processed Frame", processed_img)

    scale_percent = 50
    width = int(original_frame.shape[1] * scale_percent / 100)
    height = int(original_frame.shape[0] * scale_percent / 100)
    dim = (width, height)
    resized_frame = cv2.resize(original_frame, dim, interpolation=cv2.INTER_AREA)
    
    show_hsv(processed_img)
    cv2.imshow("Original Frame",resized_frame)

def show_hsv(processed_img):
    height, width, _ = processed_img.shape
    start_row = int(height * 0.5)
    end_row = height

    center_col = width // 2
    
    hsv_value = processed_img[start_row, center_col]
    h, s, v = hsv_value
    # print(f"HSV: ({h}, {s}, {v})")

    cv2.line(processed_img, (0, start_row), (width, start_row), (0, 255, 0), 2)  # 绿色上边界
    cv2.line(processed_img, (0, height-1), (width, height-1), (0, 255, 0), 2)    # 绿色下边界
    cv2.line(processed_img, (center_col, start_row), (center_col, height-1), (255, 0, 0), 1)  # 中心蓝线
    cv2.circle(processed_img, (center_col, start_row), 5, (0, 0, 255), -1)

    cv2.imshow("HSV Space", processed_img)


if __name__ == "__main__":
    main()
